#include <ntddk.h>

#define DEVICE_NAME        L"\\Device\\NTMODELDRV_DIRECT"
#define SYMBOLIC_LINK_NAME L"\\??\\NTMODELDRV_DIRECT"

#define DRV_PREFIX "NTMODELDRV_DIRECT:"

NTSTATUS
CompleteIrp(PIRP Irp, NTSTATUS Status, ULONG_PTR Information);

NTSTATUS
NTModelDrvDirectCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp);

NTSTATUS
NTModelDrvDirectCreateRead(PDEVICE_OBJECT DeviceObject, PIRP Irp);

NTSTATUS
NTModelDrvDirectCreateWrite(PDEVICE_OBJECT DeviceObject, PIRP Irp);

VOID
NTModelDrvDirectUnload(PDRIVER_OBJECT DriverObject);

NTSTATUS
DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegisterPath)
{
	NTSTATUS Status            = STATUS_SUCCESS;	
	DriverObject->DriverUnload = NTModelDrvDirectUnload;

	UNICODE_STRING usDeviceName       = { 0 };
	UNICODE_STRING usSymbolicLinkName = { 0 };
	PDEVICE_OBJECT DeviceObject       = NULL;

	do {
		
		RtlInitUnicodeString(&usDeviceName, DEVICE_NAME);
		RtlInitUnicodeString(&usSymbolicLinkName, SYMBOLIC_LINK_NAME);

		Status = IoCreateDevice(DriverObject, 0, &usDeviceName, FILE_DEVICE_UNKNOWN, 0, TRUE, &DeviceObject);
		if (NT_SUCCESS(Status) == FALSE) {

			KdPrintEx((77, 0, DRV_PREFIX" IoCreateDevice Failed: %x\n", Status));
			break;
		}
		DeviceObject->Flags |= DO_DIRECT_IO;

		Status = IoCreateSymbolicLink(&usSymbolicLinkName, &usDeviceName);
		if (NT_SUCCESS(Status) == FALSE) {

			KdPrintEx((77, 0, DRV_PREFIX" IoCreateSymbolicLink Failed: %x\n", Status));
			break;
		}

		DriverObject->MajorFunction[IRP_MJ_CREATE] = NTModelDrvDirectCreateClose;
		DriverObject->MajorFunction[IRP_MJ_CLOSE]  = NTModelDrvDirectCreateClose;
		DriverObject->MajorFunction[IRP_MJ_READ]   = NTModelDrvDirectCreateRead;
		DriverObject->MajorFunction[IRP_MJ_WRITE]  = NTModelDrvDirectCreateWrite;

		KdPrintEx((77, 0, DRV_PREFIX" NTModelDrv loaded successfully\n"));

	} while (FALSE);

	if (DeviceObject != NULL) {

		if (NT_SUCCESS(Status) == FALSE) {

			IoDeleteDevice(DeviceObject);
		}
	}

	return Status; 
}

VOID
NTModelDrvDirectUnload(PDRIVER_OBJECT DriverObject)
{
	UNICODE_STRING usSymbolicLinkName = { 0 };
	RtlInitUnicodeString(&usSymbolicLinkName, SYMBOLIC_LINK_NAME);

	IoDeleteSymbolicLink(&usSymbolicLinkName);
	IoDeleteDevice(DriverObject->DeviceObject);

	KdPrintEx((77, 0, DRV_PREFIX" NTModelDrv uninstalled successfully\n"));
}

NTSTATUS
CompleteIrp(PIRP Irp, NTSTATUS Status, ULONG_PTR Information)
{
	Irp->IoStatus.Status      = Status;
	Irp->IoStatus.Information = Information;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return Status;
}

NTSTATUS
NTModelDrvDirectCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	return CompleteIrp(Irp, STATUS_SUCCESS, 0);
}

NTSTATUS
NTModelDrvDirectCreateRead(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	PVOID Buffer = MmGetSystemAddressForMdlSafe(Irp->MdlAddress, NormalPagePriority);
	if (Buffer == NULL) {

		KdPrintEx((77, 0, DRV_PREFIX" MmGetSystemAddressForMdlSafe Failed\n"));

		return CompleteIrp(Irp, STATUS_INSUFFICIENT_RESOURCES, 0);
	}

	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	ULONG_PTR Len            = Stack->Parameters.Read.Length;

	WCHAR szText[] = L"Hello R3, I come from R0.";

	ULONG_PTR ReturnLen = Len < sizeof(szText) ? Len : sizeof(szText);

	RtlCopyMemory(Buffer, szText, ReturnLen);

	return CompleteIrp(Irp, STATUS_SUCCESS, ReturnLen);
}

NTSTATUS
NTModelDrvDirectCreateWrite(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	PVOID Buffer = MmGetSystemAddressForMdlSafe(Irp->MdlAddress, NormalPagePriority);
	if (Buffer == NULL) {

		KdPrintEx((77, 0, DRV_PREFIX" MmGetSystemAddressForMdlSafe Failed\n"));

		return CompleteIrp(Irp, STATUS_INSUFFICIENT_RESOURCES, 0);
	}

	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	ULONG_PTR Len            = Stack->Parameters.Write.Length;

	UNICODE_STRING usOutput = { 0 };
	RtlInitUnicodeString(&usOutput, (WCHAR*)Buffer);
	KdPrintEx((77, 0, DRV_PREFIX" %wZ\n", usOutput));

	return CompleteIrp(Irp, STATUS_SUCCESS, Len);
}